# JWT 认证工具类

from datetime import datetime, timedelta, timezone
from typing import Dict, Optional, Any
from jose import jwt, JWTError
from jose.exceptions import ExpiredSignatureError

from src.config.settings import get_settings


class JWTAuth:
    """JWT认证工具类"""

    def __init__(self):
        self.config = get_settings()
        # Refresh token 配置
        self.REFRESH_TOKEN_EXPIRE_DAYS = 30  # refresh token 30天过期

    def create_access_token(
        self,
        user_id: int,
        username: str,
        expires_delta: Optional[timedelta] = None,
        extra_data: Optional[Dict[str, Any]] = None,
    ) -> str:
        """
        创建访问token

        Args:
            user_id: 用户ID
            username: 用户名
            expires_delta: 过期时间增量，如果为None则使用配置中的默认值
            extra_data: 额外的用户数据

        Returns:
            str: JWT token
        """
        if expires_delta is None:
            expires_delta = timedelta(minutes=self.config.ACCESS_TOKEN_EXPIRE_MINUTES)

        # 计算过期时间
        expire = datetime.now(timezone.utc) + expires_delta

        # 构建payload
        payload = {
            "user_id": user_id,
            "username": username,
            "exp": expire,
            "iat": datetime.now(timezone.utc),  # 签发时间
            "type": "access_token",
        }

        # 添加额外数据
        if extra_data:
            payload.update(extra_data)

        # 生成token
        token = jwt.encode(
            payload, self.config.SECRET_KEY, algorithm=self.config.ALGORITHM
        )

        return token

    def create_refresh_token(
        self,
        user_id: int,
        username: str,
        expires_delta: Optional[timedelta] = None,
        extra_data: Optional[Dict[str, Any]] = None,
    ) -> str:
        """
        创建刷新token

        Args:
            user_id: 用户ID
            username: 用户名
            expires_delta: 过期时间增量，如果为None则使用默认30天
            extra_data: 额外的用户数据

        Returns:
            str: Refresh JWT token
        """
        if expires_delta is None:
            expires_delta = timedelta(days=self.REFRESH_TOKEN_EXPIRE_DAYS)

        # 计算过期时间
        expire = datetime.now(timezone.utc) + expires_delta

        # 构建payload
        payload = {
            "user_id": user_id,
            "username": username,
            "exp": expire,
            "iat": datetime.now(timezone.utc),  # 签发时间
            "type": "refresh_token",
        }

        # 添加额外数据
        if extra_data:
            payload.update(extra_data)

        # 生成token
        token = jwt.encode(
            payload, self.config.SECRET_KEY, algorithm=self.config.ALGORITHM
        )

        return token


    def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
        """
        验证token并返回payload

        Args:
            token: JWT token

        Returns:
            Optional[Dict[str, Any]]: 如果验证成功返回payload，否则返回None
        """
        try:
            payload = jwt.decode(
                token, self.config.SECRET_KEY, algorithms=[self.config.ALGORITHM]
            )
            return payload
        except ExpiredSignatureError:
            return None
        except JWTError:
            return None

    def get_user_from_token(self, token: str) -> Optional[Dict[str, Any]]:
        """
        从token中获取用户信息

        Args:
            token: JWT token

        Returns:
            Optional[Dict[str, Any]]: 用户信息字典
        """
        payload = self.verify_token(token)
        if payload:
            return {
                "user_id": payload.get("user_id"),
                "username": payload.get("username"),
                "exp": payload.get("exp"),
                "iat": payload.get("iat"),
            }
        return None

    def is_token_expired(self, token: str) -> bool:
        """
        检查token是否过期

        Args:
            token: JWT token

        Returns:
            bool: True表示已过期，False表示未过期
        """
        try:
            payload = jwt.decode(
                token, self.config.SECRET_KEY, algorithms=[self.config.ALGORITHM]
            )
            exp = payload.get("exp")
            if exp:
                return datetime.fromtimestamp(exp, tz=timezone.utc) < datetime.now(
                    timezone.utc
                )
            return True
        except ExpiredSignatureError:
            return True
        except JWTError:
            return True

    def refresh_access_token_from_refresh_token(
        self, refresh_token: str, expires_delta: Optional[timedelta] = None
    ) -> Optional[str]:
        """
        使用refresh token生成新的access token

        Args:
            refresh_token: Refresh JWT token
            expires_delta: 新access token的过期时间增量

        Returns:
            Optional[str]: 新的access token，如果refresh token无效则返回None
        """
        payload = self.verify_token(refresh_token)
        if payload and payload.get("type") == "refresh_token":
            user_id = payload.get("user_id")
            username = payload.get("username")

            if user_id and username:
                # 移除时间相关字段，保留其他用户数据
                extra_data = {
                    k: v
                    for k, v in payload.items()
                    if k not in ["user_id", "username", "exp", "iat", "type"]
                }

                return self.create_access_token(
                    user_id=user_id,
                    username=username,
                    expires_delta=expires_delta,
                    extra_data=extra_data if extra_data else None,
                )
        return None

    def create_token_pair(
        self,
        user_id: int,
        username: str,
        access_expires_delta: Optional[timedelta] = None,
        refresh_expires_delta: Optional[timedelta] = None,
        extra_data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, str]:
        """
        创建access token和refresh token对

        Args:
            user_id: 用户ID
            username: 用户名
            access_expires_delta: access token过期时间增量
            refresh_expires_delta: refresh token过期时间增量
            extra_data: 额外的用户数据

        Returns:
            Dict[str, str]: 包含access_token和refresh_token的字典
        """
        access_token = self.create_access_token(
            user_id=user_id,
            username=username,
            expires_delta=access_expires_delta,
            extra_data=extra_data,
        )
        
        refresh_token = self.create_refresh_token(
            user_id=user_id,
            username=username,
            expires_delta=refresh_expires_delta,
            extra_data=extra_data,
        )
        
        return {
            "access_token": access_token,
            "refresh_token": refresh_token,
        }


# 全局JWT认证实例
jwt_auth = JWTAuth()


def create_token_for_user(user_id: int, username: str, **kwargs) -> str:
    """
    便捷函数：为用户创建token

    Args:
        user_id: 用户ID
        username: 用户名
        **kwargs: 其他参数传递给create_access_token

    Returns:
        str: JWT token
    """
    return jwt_auth.create_access_token(user_id, username, **kwargs)


def verify_user_token(token: str) -> Optional[Dict[str, Any]]:
    """
    便捷函数：验证用户token

    Args:
        token: JWT token

    Returns:
        Optional[Dict[str, Any]]: 用户信息或None
    """
    return jwt_auth.get_user_from_token(token)
